// ROS
#include <rclcpp/rclcpp.hpp>

// MoveIt Servo
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

// msg
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>

// tf
#include <tf2_ros/buffer.h>
#include <tf2_ros/transform_listener.h>

// std
#include <algorithm>
#include <atomic>
#include <cmath>
#include <memory>
#include <string>

using namespace std::chrono_literals;

static const rclcpp::Logger LOGGER =
    rclcpp::get_logger("moveit2_tutorials.servo_node.cpp");

// === 全局 ===
rclcpp::Node::SharedPtr node_;
rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr wrench_sub_;
rclcpp::TimerBase::SharedPtr zero_pub_timer_;

// 可调参数（--ros-args -p xxx:=yyy）
double      g_target_force    = 0.8;            // 目标Z轴力 N
double      g_kp              = 0.10;           // 比例系数
double      g_vmax            = 0.05;           // 线速度限幅 m/s
double      g_deadband        = 0.02;           // 死区 N
double      g_stop_threshold  = 1.0;            // 超限锁存阈值 N（Fz>此值即锁存停机）
double      g_zero_pub_rate_hz= 50.0;           // 锁存后零速发布频率 Hz
std::string g_cmd_frame       = "base_link";    // 速度参考坐标系
std::string g_twist_topic     = "servo_node/delta_twist_cmds";
std::string g_wrench_topic    = "/wrench_comp";

// 锁存标志：一旦为 true，将持续发布 0 速度直到节点重启
std::atomic_bool g_stop_latched{false};

// 发布零速度（用于锁存后持续输出0）
inline void publishZeroTwist()
{
  if (!twist_cmd_pub_) return;
  auto out = std::make_unique<geometry_msgs::msg::TwistStamped>();
  out->header.stamp = node_->now();
  out->header.frame_id = g_cmd_frame;
  // 线速度与角速度均为 0
  out->twist.linear.x = out->twist.linear.y = out->twist.linear.z = 0.0;
  out->twist.angular.x = out->twist.angular.y = out->twist.angular.z = 0.0;
  twist_cmd_pub_->publish(std::move(out));
}

void wrenchCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
{
  const double fz  = msg->wrench.force.z;

  // 若已锁存，持续输出零速度
  if (g_stop_latched.load(std::memory_order_relaxed))
  {
    publishZeroTwist();
    RCLCPP_DEBUG_THROTTLE(LOGGER, *node_->get_clock(), 1000,
                          "Latched stop active; publishing zero twist.");
    return;
  }

  // 超阈值 -> 触发锁存并立即输出0
  if (fz > g_stop_threshold)
  {
    g_stop_latched.store(true, std::memory_order_relaxed);
    publishZeroTwist();
    RCLCPP_WARN(LOGGER,
                "Force=%.3f N exceeds threshold %.3f N. Latching ZERO velocity until node restart.",
                fz, g_stop_threshold);
    return;
  }

  // 正常P控制
  const double err = g_target_force - fz; // e = 目标 - 实测
  double vx_cmd = 0.0;
  if (std::fabs(err) > g_deadband)
    vx_cmd = std::clamp(g_kp * err, -g_vmax, g_vmax);

  auto out = std::make_unique<geometry_msgs::msg::TwistStamped>();
  out->header.stamp = node_->now();
  out->header.frame_id = g_cmd_frame;
  out->twist.linear.x  = vx_cmd;
  out->twist.linear.y  = 0.0;
  out->twist.linear.z  = 0.0;
  out->twist.angular.x = 0.0;
  out->twist.angular.y = 0.0;
  out->twist.angular.z = 0.0;

  twist_cmd_pub_->publish(std::move(out));

  RCLCPP_INFO(LOGGER, "Fz=%.3f N, err=%.3f -> vx=%.3f m/s", fz, err, vx_cmd);
}

int main(int argc, char** argv)
{
  rclcpp::init(argc, argv);
  rclcpp::NodeOptions node_options;
  node_options.use_intra_process_comms(false);
  node_ = std::make_shared<rclcpp::Node>("servo_node", node_options);

  // 声明并读取参数
  g_target_force     = node_->declare_parameter<double>("target_force",     g_target_force);
  g_kp               = node_->declare_parameter<double>("kp",               g_kp);
  g_vmax             = node_->declare_parameter<double>("v_max",            g_vmax);
  g_deadband         = node_->declare_parameter<double>("deadband",         g_deadband);
  g_stop_threshold   = node_->declare_parameter<double>("stop_threshold",   g_stop_threshold);
  g_zero_pub_rate_hz = node_->declare_parameter<double>("zero_pub_rate_hz", g_zero_pub_rate_hz);
  g_cmd_frame        = node_->declare_parameter<std::string>("cmd_frame",    g_cmd_frame);
  g_twist_topic      = node_->declare_parameter<std::string>("twist_topic",  g_twist_topic);
  g_wrench_topic     = node_->declare_parameter<std::string>("wrench_topic", g_wrench_topic);

  // 等一下可视化/其他节点
  rclcpp::sleep_for(4s);

  // PlanningSceneMonitor
  auto tf_buffer = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
  auto planning_scene_monitor =
      std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
          node_, "robot_description", tf_buffer, "planning_scene_monitor");

  if (planning_scene_monitor->getPlanningScene())
  {
    planning_scene_monitor->startStateMonitor("/joint_states");
    planning_scene_monitor->setPlanningScenePublishingFrequency(25.0);
    planning_scene_monitor->startPublishingPlanningScene(
        planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
        "/moveit_servo/publish_planning_scene");
    planning_scene_monitor->startSceneMonitor();
    planning_scene_monitor->providePlanningSceneService();
  }
  else
  {
    RCLCPP_ERROR(LOGGER, "规划场景未正确配置");
    return EXIT_FAILURE;
  }

  // 发布器
  twist_cmd_pub_ = node_->create_publisher<geometry_msgs::msg::TwistStamped>(
      g_twist_topic, rclcpp::QoS(10));

  // 订阅器：传感器QoS
  wrench_sub_ = node_->create_subscription<geometry_msgs::msg::WrenchStamped>(
      g_wrench_topic, rclcpp::SensorDataQoS(), &wrenchCallback);

  // 锁存后持续发布0：定时器（默认 50Hz）
  if (g_zero_pub_rate_hz > 1e-6)
  {
    const auto period_ns =
        std::chrono::duration_cast<std::chrono::nanoseconds>(
            std::chrono::duration<double>(1.0 / g_zero_pub_rate_hz));
    zero_pub_timer_ = node_->create_wall_timer(period_ns, []() {
      if (g_stop_latched.load(std::memory_order_relaxed)) publishZeroTwist();
    });
  }

  // 初始化并启动 MoveIt Servo
  auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(node_);
  if (!servo_parameters)
  {
    RCLCPP_FATAL(LOGGER, "加载 Servo 参数失败");
    return EXIT_FAILURE;
  }
  auto servo = std::make_unique<moveit_servo::Servo>(node_, servo_parameters, planning_scene_monitor);
  servo->start();

  // 执行器
  rclcpp::executors::MultiThreadedExecutor exec;
  exec.add_node(node_);
  exec.spin();

  rclcpp::shutdown();
  return 0;
}
